# RBC教材第四章可变劳动的例子
# 该算法有几个预备步骤：
# 1. 构造贝尔曼方程

rm(list = ls())
library(rootSolve)
library(signal)
# 参数设置
deltax <- 0.1
thetax <- 0.36
betax <- 0.98
A <- 0.5

# 数值计算设置
crit <- 1
tol <- 0.001
eps <- 0.02

# 格点和值函数初值设置和下一期最优资本存量内存分配
k0 <- seq(0.06,10,length.out = 80)
aopt <- rep(0,length(k0))
v <- seq(-4,9,length.out = length(k0))

j <- 0
while (crit > tol) {
  j <- j + 1
  print(crit)
  vold <- v
  m <- 0
  for (i in 1:length(k0)) {
    ainit <- k0[i]
    m <- m + 1
    v0 <- -1e10

    pos <- decise_range(k0, bellman_ch4, m0 = m,
                            a0 = ainit, a1 = k0[m], avct = k0,
                        vold = vold, thetax = thetax, deltax = deltax,
                        betax = betax, A = A)

    # 考虑边界点的最大值搜索
    if (pos$ax == pos$bx){
      pos$bx <- pos$ax + eps * (k0[2] - k0[1])
      if (interp1(k0, vold, pos$bx) < interp1(k0, vold, pos$ax)){
        aopt[i] <- k0[1]
      }else{
        aopt[i] <- GoldenSectionMax(bellman_ch4, ay = pos$ax, by = pos$bx,
                                          cy = pos$cx, a0 = ainit,
                                          betax = betax, avct = k0, vold = vold,
                                    thetax = thetax, deltax = deltax, A = A)
      }
    }else if (pos$bx == pos$cx){
      pos$bx <- pos$cx - eps * (k0[length(k0)] - k0[length(k0)-1])
      if (interp1(k0, vold,pos$bx) < interp1(k0, vold,pos$cx)){
        aopt[i] <- k0[length(k0)]
      }else {
        aopt[i] <- GoldenSectionMax(bellman_ch4, ay = pos$ax, by = pos$bx,
                                    cy = pos$cx, a0 = ainit,
                                    betax = betax, avct = k0, vold = vold,
                                    thetax = thetax, deltax = deltax, A = A)
      }
    }else {
      aopt[i] <- GoldenSectionMax(bellman_ch4, ay = pos$ax, by = pos$bx,
                                        cy = pos$cx, a0 = ainit,
                                        betax = betax, avct = k0, vold = vold,
                                        thetax = thetax, deltax = deltax, A = A)
    }
    v[i] <- bellman_ch4(a0 = k0[i], a1 = aopt[i], A = A,
                               betax = betax, thetax = thetax, deltax = deltax,
                               avct = k0, vold = vold)
  }
  crit <- abs(vold-v) %>% mean() %>% max()
}

# 获得劳动力
chfun <- function(h,a0,a1,A,thetax,deltax){
  c0 <- a0^thetax*h^(1-thetax) - a1 + (1-deltax)*a0
  A*c0/(1-h)-(1-thetax)*a0^thetax*h^(-thetax)
}

hh <- numeric(80)
for (i in 1:80) {
  hh[i] <- uniroot.all(chfun, interval = c(0,1), a0 = k0[i], a1 = aopt[i], A = A,
                       thetax = thetax, deltax = deltax)
}
# 获得消费
c0 <- k0^thetax*hh^(1-thetax) - aopt + (1-deltax)*k0


picdata <- cbind(k0,aopt,hh,c0,v) %>% as.data.frame()
ggplot(picdata, aes(x = k0, y = aopt)) + geom_line() +
  geom_line(aes(y = k0), color = 'red')

